# Copyright (c) 2022-2023, NVIDIA Corporation & Affiliates. All rights reserved. 
# 
# This work is made available under the Nvidia Source Code License-NC. 
# To view a copy of this license, visit 
# https://github.com/NVlabs/FB-BEV/blob/main/LICENSE

import datetime

from typing import Dict
import numpy as np
import os.path as osp
from mmcv.runner.dist_utils import master_only
from mmcv.runner.hooks import HOOKS, Hook
from mmcv.runner.hooks.logger.base import LoggerHook
from urllib import request, parse
import json
from urllib.error import HTTPError, URLError
import socket
from urllib import request, parse, error as urlerror
import io
import io
import json
import logging
import re
import socket
import time
from collections import defaultdict, deque
from typing import Deque, Dict, Iterable, List, Optional, Pattern
from urllib import error as urlerror, parse, request

import matplotlib.pyplot as plt
import numpy as np
import torch
import json
import logging
import re
import socket
import time
from collections import defaultdict, deque
from typing import Deque, Dict, Iterable, List, Pattern
from urllib import error as urlerror, parse, request

import matplotlib.pyplot as plt
import numpy as np
import torch
from mmcv.utils import is_list_of, is_tuple_of
from mmcv.runner import master_only

@HOOKS.register_module()
class MyGradChecker(Hook):
    def after_train_iter(self, runner):
        max_key = None
        max_val = -1e5
        min_key = None
        min_val = 1e5
        for key, val in runner.model.named_parameters():

            if val.grad.max() > max_val:
                max_val = val.grad.max()
                max_key = key
            if val.grad.min() < min_val:
                min_val = val.grad.min() 
                min_key = key
            if val.grad == None and val.requires_grad:

            
                print('WARNNING: {key}\'s parameters are not be used!!!!'.format(key=key))
        print('max grd', max_key, '  ', max_val)
        print('min grad', min_key, '  ', min_val)













@HOOKS.register_module()
class MyWechatLoggerHook(LoggerHook):
    """
    通过 Miaotixing 服务，将训练/验证指标以中文格式推送到微信。

    参数:
        interval (int): 每隔多少次迭代(训练)或多少个 epoch(验证)发送一次。
        ignore_last (bool): 是否忽略最后不足间隔的小批次。
        reset_flag (bool): 是否在日志后清空缓存。
        by_epoch (bool): 是否按 epoch 统计。
        allowed_subkeys (list[str] | None): 过滤要发送的字段，None 表示发送所有字段。
        miao_code (str): Miaotixing 的触发码。

    用法示例:
        在配置文件中添加到 `log_config.hooks`:

        log_config = dict(
            interval=50,
            hooks=[
                dict(
                    type='MyWechatLoggerHook',
                    interval=10,       # 每 10 次迭代发送一次
                    by_epoch=False,    # 按迭代发送
                    allowed_subkeys=['loss'],
                    miao_code='你的miao_code'
                ),
            ]
        )

        如果直接通过 Runner 初始化：

        runner = EpochBasedRunner(
            model=model,
            optimizer=optimizer,
            work_dir=work_dir,
            logger=logger,
            log_config=dict(
                interval=50,
                hooks=[
                    dict(
                        type='MyWechatLoggerHook',
                        interval=20,
                        by_epoch=False,
                        miao_code='your_code_here'
                    )
                ]
            )
        )
    """
    def __init__(
        self,
        interval: int = 10,
        ignore_last: bool = True,
        reset_flag: bool = False,
        commit: bool = True,
        by_epoch: bool = True,
        allowed_subkeys=None,
        miao_code: str = 'xxxxx'
    ):
        super().__init__(interval, ignore_last, reset_flag, by_epoch)
        self.miao_code = miao_code
        self.allowed_subkeys = allowed_subkeys
        self.active = True

    @master_only
    def log(self, runner):
        if not self.active:
            return
        mode = self.get_mode(runner)
        tags = self.get_loggable_tags(runner)

        # 训练阶段：每隔 interval 次迭代发送
        if mode == 'train':
            loss = tags.get('train/loss')
            if loss is None or np.isnan(loss) or np.isinf(loss):
                runner.logger.warning(
                    f"实验 {runner.meta.get('exp_name', '')} 出现无效损失：{loss}，已停止推送通知。"
                )
                self.active = False
                return
            if not self.by_epoch and (runner.iter + 1) % self.interval == 0:
                text = self.get_table_text(runner, tags)
                self._send(runner, text)
            return

        # 验证阶段：每隔 interval 个 epoch 发送
        if mode == 'val':
            if self.by_epoch and (runner.epoch + 1) % self.interval == 0:
                text = self.get_table_text(runner, tags)
                self._send(runner, text)
            return

        # 其他模式（不应出现）
        raise RuntimeError(f"未知运行模式: {mode}")

    @master_only
    def after_run(self, runner):
        text = f"实验 {runner.meta.get('exp_name', '')} 完成！"
        self._send(runner, text)

    def get_table_text(self, runner, tags):
        lines = [f"实验名称：{runner.meta.get('exp_name', '')}"]
        if self.by_epoch:
            lines.append(f"当前 Epoch：{runner.epoch + 1}")
        else:
            lines.append(f"当前 Iter：{runner.iter + 1}")
        for name, value in tags.items():
            if self.allowed_subkeys is None or any(key in name for key in self.allowed_subkeys):
                lines.append(f"{name}：{value}")
        return "\n".join(lines)

    def _send(self, runner, text):
        if not text:
            return
        url = "http://miaotixing.com/trigger"
        params = {'id': self.miao_code, 'text': text, 'type': 'json'}
        try:
            full_url = f"{url}?{parse.urlencode(params)}"
            resp = request.urlopen(full_url, timeout=5)
            runner.logger.info("MiaoTiXing：消息发送成功。")
            data = json.load(resp)
            runner.logger.debug(f"返回内容：{data}")
        except HTTPError as e:
            runner.logger.error(f"MiaoTiXing HTTP 错误：{e.code} - {e.reason}")
        except URLError as e:
            if isinstance(e.reason, socket.timeout):
                runner.logger.error(f"MiaoTiXing：访问超时，URL {url}")
            else:
                runner.logger.error(f"MiaoTiXing URLError：{e.reason}")
        except Exception as e:
            runner.logger.error(f"MiaoTiXing 未知错误：{e}")